/*
 * XInput1_3-dev - Replacement of xinput library to help with creating
 *                 drivers/debugging application
 *
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 3.0 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library.
 */

#include "stdafx.h"

#include "DeviceInfo.h"
#include "DeviceList.h"
#include "DeviceEnum.h"
#include "DriverComm.h"
#include "Protocol.h"
#include "Utilities.h"

//#include "WinDDK\hidsdi.h"

PSP_DEVICE_INTERFACE_DETAIL_DATA_W g_pDetailBuffer;
DWORD g_dwDetailBufferSize;

HRESULT GetDeviceList(HDEVINFO *phDevInfo)
{
	*phDevInfo = (HDEVINFO)-1;
	*phDevInfo = DriverComm::SetupDiWrapper::GetClassDevs(&Protocol::XUSB_INTERFACE_CLASS_GUID, 0, 0, 0x12u);

	if ( *phDevInfo == (HDEVINFO)-1 )
		return E_FAIL;

	return S_OK;
}

HRESULT GetDeviceInterfaceData(HDEVINFO DeviceInfoSet, DWORD MemberIndex, PSP_DEVICE_INTERFACE_DATA DeviceInterfaceData)
{
	DeviceInterfaceData->cbSize = 0;
	DeviceInterfaceData->InterfaceClassGuid.Data1 = 0;
	DeviceInterfaceData->InterfaceClassGuid.Data2 = 0;
	*(DWORD*)&DeviceInterfaceData->InterfaceClassGuid.Data4[0] = 0;
	*(DWORD*)&DeviceInterfaceData->InterfaceClassGuid.Data4[4] = 0;
	DeviceInterfaceData->Flags = 0;
	DeviceInterfaceData->Reserved = 0;
	DeviceInterfaceData->cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);

	if ( DriverComm::SetupDiWrapper::EnumDeviceInterfaces(DeviceInfoSet, 0, &Protocol::XUSB_INTERFACE_CLASS_GUID, MemberIndex, DeviceInterfaceData) )
		return S_OK;

	return E_FAIL;
}

signed int __stdcall GetDeviceInterfaceDetail(HDEVINFO DeviceInfoSet, PSP_DEVICE_INTERFACE_DATA pDeviceInterfaceData, PSP_DEVICE_INTERFACE_DETAIL_DATA_W *ppDeviceInterfaceDetail)
{
	SIZE_T RequiredSize = g_dwDetailBufferSize;
	*ppDeviceInterfaceDetail = g_pDetailBuffer;

	if ( *ppDeviceInterfaceDetail )
	{
		memset(*ppDeviceInterfaceDetail, 0, g_dwDetailBufferSize);
		(*ppDeviceInterfaceDetail)->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);
	}

	if ( DriverComm::SetupDiWrapper::GetDeviceInterfaceDetail(DeviceInfoSet, pDeviceInterfaceData, *ppDeviceInterfaceDetail, g_dwDetailBufferSize, &RequiredSize, 0) )
		return S_OK;

	*ppDeviceInterfaceDetail = 0;

	if ( GetLastError() != 122 )
	{
		return E_FAIL;
	}

	if ( g_pDetailBuffer )
	{
		LocalFree(g_pDetailBuffer);
		g_pDetailBuffer = 0;
		g_dwDetailBufferSize = 0;
	}

	g_pDetailBuffer = (PSP_DEVICE_INTERFACE_DETAIL_DATA_W)LocalAlloc(0x40u, RequiredSize);

	if ( !g_pDetailBuffer )
		return E_OUTOFMEMORY;

	g_dwDetailBufferSize = RequiredSize;
	*ppDeviceInterfaceDetail = g_pDetailBuffer;

	(*ppDeviceInterfaceDetail)->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);

	if ( DriverComm::SetupDiWrapper::GetDeviceInterfaceDetail(DeviceInfoSet, pDeviceInterfaceData, *ppDeviceInterfaceDetail, g_dwDetailBufferSize, 0, 0) )
	{
		return S_OK;
	}

	*ppDeviceInterfaceDetail = 0;
	return E_FAIL;
}

HRESULT ProcessEnumeratedDevice(PDEVICE_INFO pDeviceInfo)
{
	HRESULT hResult = DeviceInfo::IsDeviceActive(pDeviceInfo);
	if ( FAILED(hResult) || hResult == S_FALSE )
		return hResult;

	hResult = DeviceList::IsDevicePresent(pDeviceInfo);
	if ( FAILED(hResult) )
		return hResult;

	if ( hResult == S_OK )
		return S_FALSE;

	BYTE bPort = 0xFF;
	BYTE bLEDState = 0;

	if ( (DriverComm::GetLEDState(pDeviceInfo, &bLEDState) & 0x80000000) == 0 )
	{
		if ( bLEDState < Protocol::LEDState::XINPUT_LED_MAX )
		{
			bPort = Protocol::LEDState::XINPUT_LED_TO_PORT_MAP[bLEDState];
			if ( bPort != 0xFF )
			{
				hResult = DeviceList::IsDeviceOnPort(bPort);
				if ( FAILED(hResult) )
					return hResult;

				if ( hResult != S_FALSE )
				{
					bPort = 0xFF;
					DriverComm::SendLEDState(pDeviceInfo, Protocol::LEDState::XINPUT_LED_FLIPFLOP);
				}
			}
		}
	}

	for (DWORD i = 0; i < XUSER_MAX_COUNT && bPort == 255; ++i )
	{
		hResult = DeviceList::IsDeviceOnPort(i);
		if ( FAILED(hResult) )
			return hResult;

		if ( hResult == S_FALSE )
			bPort = (BYTE)i;
	}

	if ( bPort >= XUSER_MAX_COUNT )
	{
		return S_FALSE;
	}

	hResult = DeviceList::SetDeviceOnPort(bPort, pDeviceInfo);
	if ( FAILED(hResult) )
		return hResult;

	if ( !Utilities::IsSettingSet(1) )
		DriverComm::SendLEDState(pDeviceInfo, Protocol::LEDState::XINPUT_PORT_TO_LED_MAP[bPort % 4]);

	return S_OK;
}

HRESULT EnumerateDevicesOnDeviceInterface(HANDLE hSourceHandle, wchar_t *Str)
{
	PDEVICE_INFO pDeviceInfo = 0;
	INTERFACE_INFORMATION InterfaceInformation = {0};

	HRESULT hResult = DriverComm::GetDeviceInfoFromInterface(hSourceHandle, &InterfaceInformation);
	if ( (hResult & 0x80000000) == 0 && !(InterfaceInformation.dwUsers & 0x80) )
	{
		for (DWORD i = 0; i < InterfaceInformation.dwUsers; ++i )
		{
			if ( !pDeviceInfo )
			{
				pDeviceInfo = DeviceInfo::Create(hSourceHandle, Str);
				if ( !pDeviceInfo )
					return E_OUTOFMEMORY;
			}
			else
			{
				DeviceInfo::Recycle(pDeviceInfo);
			}

			pDeviceInfo->bPort = i;
			pDeviceInfo->Dummy78 = InterfaceInformation.wVID;
			pDeviceInfo->Dummy80 = InterfaceInformation.wPID;
			pDeviceInfo->wVersion = InterfaceInformation.wVersion;

			hResult = DriverComm::GetLatestDeviceInfo( pDeviceInfo );
			if ( (hResult & 0x80000000) == 0 )
			{
				hResult = ProcessEnumeratedDevice(pDeviceInfo);
				if ( !hResult )
					pDeviceInfo = 0;
			}
		}

		if ( pDeviceInfo )
		{
			DeviceInfo::Destroy(pDeviceInfo);
			pDeviceInfo = 0;
		}
	}

	return hResult;
}

HRESULT EnumerateXInputDevices()
{
	HDEVINFO DeviceInfoSet = (HDEVINFO)-1;
	DWORD MemberIndex = 0;
	HANDLE hSourceHandle = INVALID_HANDLE_VALUE;
	bool bContinue = 1;
	
	SP_DEVICE_INTERFACE_DATA DeviceInterfaceData;
	PSP_DEVICE_INTERFACE_DETAIL_DATA_W pDeviceInterfaceDetail = 0;

	HRESULT hResult = GetDeviceList(&DeviceInfoSet);
	if ( FAILED(hResult) )
		bContinue = false;

	while ( bContinue )
	{
		pDeviceInterfaceDetail = 0;
		if ( SUCCEEDED(GetDeviceInterfaceData(DeviceInfoSet, MemberIndex, &DeviceInterfaceData)) )
		{
			MemberIndex++;
			hResult = GetDeviceInterfaceDetail(DeviceInfoSet, &DeviceInterfaceData, &pDeviceInterfaceDetail);
			if ( SUCCEEDED(hResult) )
			{
				hSourceHandle = CreateFileW(pDeviceInterfaceDetail->DevicePath, 0xC0000000u, 3u, 0, 3u, 0x80u, 0);
				if ( hSourceHandle != (HANDLE)-1 )
				{
					hResult = EnumerateDevicesOnDeviceInterface(hSourceHandle, pDeviceInterfaceDetail->DevicePath);
					if ( FAILED(hResult) )
						hResult = S_OK;

					CloseHandle(hSourceHandle);
					hSourceHandle = INVALID_HANDLE_VALUE;
				}
			}
		}
		else
		{
			hResult = S_OK;
			bContinue = false;
		}
	}

	if ( DeviceInfoSet != (HDEVINFO)-1 )
	{
		DriverComm::SetupDiWrapper::DestroyDeviceInfoList(DeviceInfoSet);
		DeviceInfoSet = (HDEVINFO)-1;
	}

	return hResult;
}

HRESULT DeviceEnum::Close(void)
{
	if ( g_pDetailBuffer )
	{
		LocalFree(g_pDetailBuffer);
		g_pDetailBuffer = 0;
	}
	g_dwDetailBufferSize = 0;
	return S_OK;;
}

HRESULT DeviceEnum::GetDeviceOnPort(DWORD dwPort,PDEVICE_INFO *ppDeviceInfo,bool bEnumerate)
{
	HRESULT hResult = DeviceList::GetDeviceOnPort(dwPort, ppDeviceInfo);
	if( SUCCEEDED(hResult) )
	{
		if( !*ppDeviceInfo && bEnumerate )
		{
			hResult = EnumerateXInputDevices();
			if( SUCCEEDED(hResult) )
			{
				hResult = DeviceList::GetDeviceOnPort(dwPort, ppDeviceInfo);
			}
		}
	}

	return hResult;
}

HRESULT DeviceEnum::Initialize(void)
{
	g_pDetailBuffer = 0;
	g_dwDetailBufferSize = 0;
	return S_OK;
}


QuickDriverEnum::QuickDriverEnum(void)
{
	this->m_hDevInfo = (HDEVINFO)-1;
	this->m_dwMemberIndex = 0;
	this->m_pDeviceInterfaceDetail = NULL;
	this->m_dwDeviceInterfaceDetailDataSize = 0;
}

QuickDriverEnum::~QuickDriverEnum(void)
{
	if ( this->m_hDevInfo != (HDEVINFO)-1 )
	{
		DriverComm::SetupDiWrapper::DestroyDeviceInfoList(this->m_hDevInfo);
		this->m_hDevInfo = (HDEVINFO)-1;
		this->m_dwMemberIndex = 0;
	}

	if ( this->m_pDeviceInterfaceDetail )
	{
		operator delete(this->m_pDeviceInterfaceDetail);
		this->m_pDeviceInterfaceDetail = NULL;
		this->m_dwDeviceInterfaceDetailDataSize = 0;
	}
}

bool QuickDriverEnum::GetNext(HANDLE *phDevice)
{
	DWORD RequiredSize;
	SP_DEVICE_INTERFACE_DATA DeviceInterfaceData;

	if ( !phDevice )
		return false;

	*phDevice = INVALID_HANDLE_VALUE;

	if ( this->m_hDevInfo == (HDEVINFO)-1 || !this->m_pDeviceInterfaceDetail )
		return false;

	DeviceInterfaceData.InterfaceClassGuid.Data1 = 0;
	DeviceInterfaceData.InterfaceClassGuid.Data2 = 0;
	*(DWORD *)&DeviceInterfaceData.InterfaceClassGuid.Data4[0] = 0;
	*(DWORD *)&DeviceInterfaceData.InterfaceClassGuid.Data4[4] = 0;
	DeviceInterfaceData.Flags = 0;
	DeviceInterfaceData.Reserved = 0;
	DeviceInterfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);

	if ( !DriverComm::SetupDiWrapper::EnumDeviceInterfaces(this->m_hDevInfo, 0, &Protocol::XUSB_INTERFACE_CLASS_GUID, this->m_dwMemberIndex, &DeviceInterfaceData) )
		return false;

	this->m_dwMemberIndex++;
	this->m_pDeviceInterfaceDetail->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);
	RequiredSize = 0;

	if ( !DriverComm::SetupDiWrapper::GetDeviceInterfaceDetail(this->m_hDevInfo, &DeviceInterfaceData, this->m_pDeviceInterfaceDetail, this->m_dwDeviceInterfaceDetailDataSize, &RequiredSize, 0) )
	{
		if ( GetLastError() != ERROR_INSUFFICIENT_BUFFER )
			return false;

		operator delete(this->m_pDeviceInterfaceDetail);
		this->m_dwDeviceInterfaceDetailDataSize = RequiredSize;
		this->m_pDeviceInterfaceDetail = (PSP_DEVICE_INTERFACE_DETAIL_DATA_W)(operator new(this->m_dwDeviceInterfaceDetailDataSize));
		if ( !this->m_pDeviceInterfaceDetail )
		{
			this->m_dwDeviceInterfaceDetailDataSize;
			return false;
		}

		this->m_pDeviceInterfaceDetail->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);
		if( !DriverComm::SetupDiWrapper::GetDeviceInterfaceDetail(this->m_hDevInfo, &DeviceInterfaceData, this->m_pDeviceInterfaceDetail, this->m_dwDeviceInterfaceDetailDataSize, 0, 0) )
		{
			return false;
		}
	}

	*phDevice = CreateFileW(this->m_pDeviceInterfaceDetail->DevicePath, 0xC0000000, 3, 0, 3, 0x80, 0);
	return true;
}

bool QuickDriverEnum::Restart(void)
{
	if ( this->m_hDevInfo != (HDEVINFO)-1 )
	{
		DriverComm::SetupDiWrapper::DestroyDeviceInfoList(this->m_hDevInfo);
		this->m_hDevInfo = (HDEVINFO)-1;
	}

	this->m_dwMemberIndex = 0;
	if ( !this->m_pDeviceInterfaceDetail )
	{
		this->m_dwDeviceInterfaceDetailDataSize = 262;
		this->m_pDeviceInterfaceDetail = (PSP_DEVICE_INTERFACE_DETAIL_DATA_W)operator new(this->m_dwDeviceInterfaceDetailDataSize);
	}
	
	if( this->m_pDeviceInterfaceDetail )
	{
		this->m_hDevInfo = DriverComm::SetupDiWrapper::GetClassDevs(&Protocol::XUSB_INTERFACE_CLASS_GUID, 0, 0, 0x12u);
		return ( this->m_hDevInfo != (HDEVINFO)-1 );
	}
	
	this->m_dwDeviceInterfaceDetailDataSize = 0;
	
	return false;
}